# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: jasonyoon
# pylint: disable=wrong-import-position,too-few-public-methods
# pylint: disable=unnecessary-comprehension,too-many-locals,broad-except
# pylint: disable=no-name-in-module
import os
from os.path import join as pjoin
import time
import json
import datetime
import tempfile

import tqdm
import fire
import numpy as np
import pandas as pd

from google.protobuf.json_format import MessageToDict

from coin.support.market_quotes.logic.market_status_util \
  import query_history_kline_list
from coin.support.market_quotes.logic.query_util import (
    _get_api_version_for_request, KlineRequest)
from xunkemgmt_client.tool.slack_noti import send_to_slack
from coin2.strategy.basis_strat.klines_pb2 import KlineProto, KlineListProto
from xunkemgmt_client.client.util.query_util import query_exchange_apis


os.environ["TZ"] = "GMT"
time.tzset()

LONGONLY = ["Gdax", "Bithumb"]


CURRENCY_MAP = {
  "Binance": "USDT",
  "Okex": "USDT",
  "Huobi": "USDT",
  "Ftx": "USD",
  "Gdax": "USD",
  "Bybit": "USDT",
  "Bithumb": "KRW"
}

EXTRA_SYMBOLS = {
  "Binance": [
    ("1000SHIB-USDT", "Futures"), ("SHIB-USDT", "Spot"),
    ("1000XEC-USDT", "Futures"), ("XEC-USDT", "Spot")],
}
MAX_CNT = {
  "Binance": 499,
}

def _get_kline_request(trading_date, product, kline_period):
  market_type, exchange, symbol = product.split(':')
  exchange_api_df = query_exchange_apis() \
    .set_index(['market_type', 'exchange', 'api_version'])
  api_version = _get_api_version_for_request(market_type, exchange, symbol)
  assert (market_type, exchange, api_version) in exchange_api_df.index
  source_exchange_api_id = exchange_api_df.loc[
      (market_type, exchange, api_version), 'exchange_api_id']
  return KlineRequest(
      market_type, exchange, symbol, kline_period, source_exchange_api_id,
      trading_date, time.time())


def _query_kline(request, start_time, end_time):
  return query_history_kline_list([request], start_time, end_time)


def query_kline_as_pb(trading_date, product,
                      start_time, end_time, kline_period='5m',
                      max_cnt=200, tolerance=2):
  request = _get_kline_request(trading_date, product, kline_period)

  # get period in secs
  if kline_period == "1m":
    period = 60
  elif kline_period == "3m":
    period = 180
  elif kline_period == "5m":
    period = 300
  elif kline_period == "15m":
      period = 900
  elif kline_period == "30m":
      period = 1800
  elif kline_period == "1h":
    period = 3600
  elif kline_period == '4h':
    period = 3600 * 4
  elif kline_period == "1d":
    period = 3600 * 24
  elif kline_period == "2d":
    period = 2 * 3600 * 24
  elif kline_period == "3d":
    period = 3 * 3600 * 24
  else:
    raise ValueError(f"invalid period: {kline_period}")
  period *= max_cnt

  klinelist = KlineListProto()
  failed = []
  query_times = []
  cur_start = start_time
  while cur_start < end_time:
    cur_end = cur_start + datetime.timedelta(seconds=period)
    query_times.append([cur_start, min(cur_end, end_time)])
    cur_start = cur_end

  for _start, _end in query_times:
    time.sleep(0.15)
    klines, error_cnt = None, 0
    while error_cnt < tolerance:
      try:
        klines = _query_kline(request, _start, _end)[0]
        assert klines is not None
        break
      except Exception as e:
        error_cnt += 1
        print(f"error: {str(e)}, product: {product}, error_cnt: {error_cnt}")
        time.sleep(1)
    if error_cnt >= tolerance:
      print(f"skip {trading_date} {product}")
      print(f"error_cnt: {error_cnt}, tolerance: {tolerance}")
      failed.append({
        "product": product,
        "kline_period": kline_period,
        "trading_date": trading_date.strftime("%Y%m%d_%H%M%S"),
        "start": _start.strftime("%Y%m%d_%H%M%S"),
        "end": _end.strftime("%Y%m%d_%H%M%S")
      })
    if klines is not None:
      for kline in klines.klines.klines:
        kline_proto = KlineProto()
        kline_proto.open = kline.open
        kline_proto.high = kline.high
        kline_proto.low = kline.low
        kline_proto.close = kline.close
        kline_proto.volume = kline.volume
        kline_proto.kline_timestamp = kline.kline_timestamp
        kline_proto.buy_volume = kline.buy_volume
        kline_proto.sell_volume = kline.sell_volume
        kline_proto.quote_volume = kline.turnover
        kline_proto.buy_quote_volume = kline.buy_turnover
        kline_proto.sell_quote_volume = kline.sell_turnover
        klinelist.klines.append(kline_proto)
  return klinelist, failed


def get_pi_path_by_exchange(market_type, exchange):
  pi_path_root = os.path.expanduser("data/coin2/product_info")
  futures_pi_name_map = {
      "Binance": ["Futures.Binance.v1.json", "Futures.Binance.v1-delivery.json"],
      "Huobi": ["Futures.Huobi.v1-linear-swap.json"],
      "Okex": ["Futures.Okex.v5-swap.json", "Futures.Okex.v5.json"],
      "Ftx": ["Futures.Ftx.v1.json"],
      "Dydx": ["Futures.Dydx.v3.json"],
      "Bybit": ["Futures.Bybit.v3.json", "Futures.Bybit.v3-linear.json", "Futures.Bybit.v3-linear-usdc.json"],
      "Bitget": ["Futures.Bitget.v1.json"],
      "Phemex": ["Futures.Phemex.v1.json"],
      "Kucoin": ["Futures.Kucoin.v1.json"],
      "Deribit": ["Futures.Deribit.v2.json"],
  }
  spot_pi_name_map = {
      "Binance": ["Spot.Binance.v1.json"],
      "Huobi": ["Spot.Huobi.v1.json"],
      "Okex": ["Spot.Okex.v5.json"],
      "Ftx": ["Spot.Ftx.v1.json"],
      "Gdax": ["Spot.Gdax.v1.json"],
      "Upbit": ["Spot.Upbit.v1.json"],
      "Bithumb": ["Spot.Bithumb.v2.json"],
      "Bitbank": ["Spot.Bitbank.v1.json"],
      "Bybit": ["Spot.Bybit.v3.json"],
      "Btcturk": ["Spot.Btcturk.v1.json"],
      'Gateio': ["Spot.Gateio.v4.json"],
      'Mexc': ["Spot.Mexc.v3.json"],
      'Kucoin': ["Spot.Kucoin.v1.json"],
      'Kraken': ["Spot.Kraken.v1.json"],
      'Korbit': ["Spot.Korbit.v1.json"],
      "Bitget": ["Spot.Bitget.v1.json"],
      "Phemex": ["Spot.Phemex.v1.json"],
      "Coingecko": ["Spot.Coingecko.v3.json"],
  }
  options_pi_name_map = {
      "Deribit": ["Options.Deribit.v2.json"],
  }
  assert market_type in ['Futures', 'Spot'], market_type

  ret = list()
  if market_type == "Futures":
    assert exchange in futures_pi_name_map, exchange
    for pi_file_name in futures_pi_name_map[exchange]:
      ret.append(pjoin(pi_path_root, pi_file_name))
  elif market_type == "Spot":
    assert exchange in spot_pi_name_map, exchange
    for pi_file_name in spot_pi_name_map[exchange]:
      ret.append(pjoin(pi_path_root, pi_file_name))
  elif market_type == "Options":
    assert exchange in options_pi_name_map, exchange
    for pi_file_name in options_pi_name_map[exchange]:
      ret.append(pjoin(pi_path_root, pi_file_name))
  return ret


def get_symbols_from_pi(market_type, exchange, append_contract_value=False, futures_suffix=".PERPETUAL"):
  pi_paths = get_pi_path_by_exchange(market_type, exchange)
  print(pi_paths)
  pi_info = list()
  for pi_path in pi_paths:
    with open(pi_path, "r", encoding="utf8") as fin:
      pi_info_each = json.loads(fin.read())["product_infos"]
      pi_info.extend(pi_info_each)
  symbols_with_pi = list()
  for _pi in pi_info:
    symbols_with_pi.append({
      'symbol': _pi["symbol"],
      'contract_value': _pi.get("contract_value", 1.0),
      'is_inverse': _pi.get("is_inverse", False),
    })
  if market_type == "Futures":
    symbols_with_pi = list(filter(lambda u: u['symbol'].endswith(futures_suffix), symbols_with_pi))
  print(f"symbols from {pi_path}, len: {len(symbols_with_pi)}, head: {symbols_with_pi[:5]}")
  if append_contract_value:
    return symbols_with_pi
  return list(map(lambda u: u['symbol'], symbols_with_pi))

def convert_klinelist(klinelist):
  return MessageToDict(klinelist,
                        including_default_value_fields=True,
                        preserving_proto_field_name=True)

def run(now=None, kline_path=None, test=False, exchange="Binance",
        use_slack=False, tolerance=2, ckpt_days=5):
  el_st = time.time()
  if now is None:
    now = datetime.datetime.now()
  elif isinstance(now, (int, str)):
    now = datetime.datetime.strptime(str(now), "%Y%m%d%H%M")
  else:
    raise ValueError(now)
  today = datetime.datetime.strptime(now.strftime("%Y%m%d"), "%Y%m%d") \
    - datetime.timedelta(seconds=1)
  ckpt = today - datetime.timedelta(days=ckpt_days)
  start = ckpt - datetime.timedelta(days=21 - ckpt_days)
  print(f"start: {start}, ckpt: {ckpt}, end: {now}")

  currency = CURRENCY_MAP[exchange]
  if exchange not in LONGONLY:
    fut_symbols = [symbol.split(".")[0] for symbol in
                   get_symbols_from_pi("Futures", exchange)]
  else:
    fut_symbols = None
  spot_symbols = get_symbols_from_pi("Spot", exchange)
  print("fut_symbols: %s" % fut_symbols)
  print("spot_symbols: %s" % spot_symbols)

  if fut_symbols:
    symbols = set(fut_symbols) & set(spot_symbols)
  else:
    symbols = set(spot_symbols)
  symbols = sorted(symbols)
  symbols = [symbol for symbol in symbols if symbol.split("-")[1] == currency]
  kline_info = {"klines": []}
  failed_info = {"failed": []}
  if test:
    np.random.seed(777)
    symbols = np.random.choice(symbols, 5)
  print(f"symbols len: {len(symbols)}, head: {symbols[:5]}")

  def add_symbol(symbol, market_type):
    _symbol = f"{symbol}.PERPETUAL" if market_type == "Futures" else symbol
    klinelist, failed = query_kline_as_pb(
        datetime.datetime.now(),
        f"{market_type}:{exchange}:{_symbol}",
        start, ckpt, kline_period="1d", tolerance=tolerance,
        max_cnt=MAX_CNT.get(exchange, 200))
    klinelist2, failed2 = query_kline_as_pb(
        datetime.datetime.now(),
        f"{market_type}:{exchange}:{_symbol}",
        ckpt, now, kline_period="5m", tolerance=tolerance,
        max_cnt=MAX_CNT.get(exchange, 200))
    klinelist = convert_klinelist(klinelist)
    klinelist2 = convert_klinelist(klinelist2)
    failed_info["failed"] += failed + failed2
    suffix = "Futures" if market_type == "Futures" else ""
    kline_info["klines"].append({
        "symbol": f"{exchange}{suffix}_{_symbol}",
        "ohlc": klinelist["klines"] + klinelist2["klines"]
    })
    # print(f"kline of {_symbol}, {market_type} fetched")

  market_types = ["Spot"]
  if exchange not in LONGONLY:
    market_types.append("Futures")

  for symbol in tqdm.tqdm(symbols):
    for market_type in market_types:
      add_symbol(symbol, market_type)
  for symbol, market_type in EXTRA_SYMBOLS.get(exchange, []):
    add_symbol(symbol, market_type)
  dt_str = now.strftime("%Y%m%d%H")
  kline_path = kline_path or \
    pjoin(tempfile.mkdtemp(prefix=f"coin_midfreq_kline_{dt_str}", dir="/tmp"),
          "kline_info.json")
  print(f"dump kline to {kline_path}")
  with open(kline_path, "w", encoding="utf8") as fout:
    fout.write(json.dumps(kline_info, indent=2))
  elapsed = time.time() - el_st
  print(f"elapsed to get kline: {elapsed:.4e} sec")
  if use_slack:
    channel = "#coin_midfreq_checker" if not test else "#stock_ops_checker_test"
    if failed_info["failed"]:
      failed_df = pd.DataFrame(failed_info["failed"])
      failed_cnt = failed_df['product'].value_counts().sort_values()
      send_to_slack('\n'.join([
                    f'Hostname: {os.environ.get("HOSTNAME", "")}',
                    f'kline_path: {kline_path}',
                    "",
                    "Failure counts",
                    failed_cnt.to_string(),
                    "",
                    "Failures",
                    failed_df.to_string()]),
                    channel,
                    'file',
                    now.strftime(f"Kline failures %Y%m%d_%H%M%S {exchange}"))


if __name__ == "__main__":
  fire.Fire()
